package com.farm2.llm.client;

import com.farm2.llm.client.impl.Farm2EmbedAliyunClientImpl;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.farm2.tools.json.FarmJsons;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class EmbedClientTest {


    public static void main(String[] args) {
        String apiKey = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxx";
        String baseUrl = "http://192.168.0.107:11434/api/embed";
        HttpClient client = HttpClient.newHttpClient();

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "nomic-embed-text");
        requestBody.put("input", "风急天高猿啸哀，渚清沙白鸟飞回，无1111111边落木萧萧下，不尽长江滚滚来");



        try {
            String requestBodyString = FarmJsons.toJson(requestBody);
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(baseUrl))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + apiKey)
                    .POST(HttpRequest.BodyPublishers.ofString(requestBodyString))
                    .build();

            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

            if (response.statusCode() == 200) {
                float[] embeddings = getEmbeddings(response.body());
                System.out.println("Response: " + embeddings.length);
            } else {
                System.out.printf("Failed to retrieve response, status code: %d, response: %s%n", response.statusCode(), response.body());
            }
        } catch (Exception e) {
            System.err.println("Error: " + e.getMessage());
        }
    }

    public static float[] getEmbeddings(String json) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {

            Map<String, Object> responseMap = objectMapper.readValue(json, new TypeReference<Map<String, Object>>() {});

            @SuppressWarnings("unchecked")
            List<List<Double>> embeddings = (List<List<Double>>) responseMap.get("embeddings");

            if (embeddings == null || embeddings.isEmpty()) {
                throw new IllegalArgumentException("Embeddings field is missing or empty in the JSON response.");
            }

            List<Double> embeddingList = embeddings.get(0);

            float[] embeddingArray = new float[embeddingList.size()];
            for (int i = 0; i < embeddingList.size(); i++) {
                embeddingArray[i] = embeddingList.get(i).floatValue();
            }

            return embeddingArray;
        } catch (Exception e) {
            throw new RuntimeException("Error parsing JSON: " + e.getMessage(), e);
        }
    }
}
